import os
import cv2
import numpy as np
from scipy.spatial.transform import Rotation as R
import sapien
from pathology.envs.pathology_env import PathologyEnv


env = PathologyEnv(headless=False)

for episode in range(2):
    obs = env.reset()
    phase = 0

    for step in range(500000):

        # if phase % 4 == 0:
        #     result = env.slide_planner.phase1(obs)
        # elif phase % 4 == 1:
        #     result = env.slide_planner.phase2(obs)
        # elif phase % 4 == 2:
        #     result = env.slide_planner.phase3(obs)
        # elif phase % 4 == 3:
        #     result = env.slide_planner.phase4(obs)
        # phase += 1

        if phase == 0:
            result = env.slide_planner.phase0_ibvs(obs)
            phase += 1
        else:
            result = env.slide_planner.phase1_ibvs(obs)
            action = result['position'][-1]
            result = {'status': 'Skip'}
            

        extra = [
                # 'slide_boundingbox',
                'camera',
                ]

        if result["status"] == "Success":
            for i in range(len(result['position'])):
                action = result['position'][i]
                obs, reward, done, info = env.step(action)
                env.render(extra)
        else:
            obs, reward, done, info = env.step(action)
            env.render(extra)
        
        if done:
            print(f'Done at step {step}')
            break

        # print("========== verbose ===========")
        # print('pose gt:', env.tool0.get_pose())
        # print('pose mplib', env.slide_planner.forward_kinematic(env.robot.get_qpos()))
        # pybullet_T = env.pybullet_planner.solve_fk_homogeneous(env.robot.get_qpos())
        # print('pose pybullet:', sapien.Pose(pybullet_T))
        # print('vel gt:', env.tool0.get_linear_velocity(), ' --- ', env.tool0.get_angular_velocity())
        # print('vel mplib:', env.slide_planner.forward_velocity(env.robot.get_qpos(), env.robot.get_qvel()))
        # pybullet_vel = env.pybullet_planner.solve_fk_velocity(env.robot.get_qpos(), env.robot.get_qvel())
        # print('vel pybullet:', pybullet_vel)

env.close()
cv2.destroyAllWindows()